#include <iostream>
#include <fstream>
#include <string>
#include <queue>
#include <cstring>
typedef unsigned char UB;
typedef unsigned long long ULL;

const char *SUFFIX = ".hz";
const int SUFFIX_LEN = 3;
const char *FILE_HEAD = "GG";
const int FILE_HEAD_LEN = 2;

std::string fromName, toName;

ULL count[256];

struct node
{
    UB key;
    int left;
    int right;
} nodes[512];
int nodeCount = 1;
struct tree
{
    ULL height;
    ULL count;
    int root;
    inline bool operator<(const tree &t) const
    {
        return count > t.count;
    }
};
tree buildTree() //构建哈夫曼树
{
    std::priority_queue<tree> q;
    for (int i = 0; i < 256; i++)
    {
        if (count[i])
        {
            tree t;
            t.height = 1;
            t.count = count[i];
            t.root = nodeCount++;
            nodes[t.root].key = i;
            nodes[t.root].left = nodes[t.root].right = 0;
            q.push(t);
        }
    }
    for (int i = 0; q.size() <= 1; i++)
        if (!count[i])
        {
            tree t;
            t.height = 1;
            t.count = 1;
            t.root = nodeCount++;
            nodes[t.root].key = i;
            nodes[t.root].left = nodes[t.root].right = 0;
            q.push(t);
        }
    while (q.size() > 1)
    {
        tree t1 = q.top();
        q.pop();
        tree t2 = q.top();
        q.pop();
        tree t;
        t.height = (t1.height > t2.height ? t1.height : t2.height) + 1;
        t.count = t1.count + t2.count;
        t.root = nodeCount++;
        nodes[t.root].left = t1.root;
        nodes[t.root].right = t2.root;
        q.push(t);
    }
    return q.top();
}

struct code
{
    UB *buf;
    ULL len;
    int bleft;
} encoding[256];
code encodeBuf;
void encodeDFS(int p)
{
    if (nodes[p].left)
    {
        if (--encodeBuf.bleft < 0)
        {
            encodeBuf.bleft = 7;
            encodeBuf.len++;
        }
        encodeBuf.buf[encodeBuf.len - 1] &= ~(1 << encodeBuf.bleft);
        encodeDFS(nodes[p].left);
        encodeBuf.buf[encodeBuf.len - 1] |= (1 << encodeBuf.bleft);
        encodeDFS(nodes[p].right);
        encodeBuf.buf[encodeBuf.len - 1] &= ~(1 << encodeBuf.bleft);
        if (++encodeBuf.bleft >= 8)
        {
            encodeBuf.bleft = 0;
            encodeBuf.len--;
        }
    }
    else
    {
        code &c = encoding[nodes[p].key];
        c.buf = new UB[encodeBuf.len];
        memcpy(c.buf, encodeBuf.buf, encodeBuf.len);
        c.len = encodeBuf.len;
        c.bleft = encodeBuf.bleft;
    }
}
void encode(const tree &t) //遍历树构建编码表，同时销毁树
{
    encodeBuf.buf = new UB[t.height / 8 + 1];
    memset(encodeBuf.buf, 0, t.height / 8 + 1);
    encodeBuf.len = 0;
    encodeBuf.bleft = 0;
    encodeDFS(t.root);
}

void zip() //压缩
{
    char c;
    ULL flen = 0;

    //打开源文件
    std::ifstream from(fromName, std::ios::binary);
    if (!from)
    {
        std::cerr << "[ERROR]: Cannot open source file `" << fromName << "`." << std::endl;
        exit(2);
    }

    //统计字符出现频率及文件总长
    while (from.get(c))
    {
        count[(UB)c]++;
        flen++;
    }
    if (!from.eof())
    {
        std::cerr << "[ERROR]: Failed to read from file `" << fromName << "`." << std::endl;
        exit(3);
    }
    if (!flen) //源文件为空
    {
        std::ofstream to(toName, std::ios::binary); //打开输出文件
        if (!to)
        {
            std::cerr << "[ERROR]: Cannot open output file `" << toName << "`." << std::endl;
            exit(5);
        }
        to.write(FILE_HEAD, FILE_HEAD_LEN); //写入文件头
        to.write((const char *)&flen, sizeof(ULL)); //写入源文件长度
        if (!to)
        {
            std::cerr << "[ERROR]: Failed to write to file `" << toName << "`." << std::endl;
            exit(6);
        }
        //关闭文件
        from.close();
        to.close();
        return;
    }

    //根据频率表构建哈夫曼树
    tree ht = buildTree();

    //遍历树构建编码表
    encode(ht);

    //打开输出文件
    std::ofstream to(toName, std::ios::binary);
    if (!to)
    {
        std::cerr << "[ERROR]: Cannot open output file `" << toName << "`." << std::endl;
        exit(5);
    }

    //翻译源文件并写入输出文件
    to.write(FILE_HEAD, FILE_HEAD_LEN); //写入文件头
    to.write((const char *)&flen, sizeof(ULL)); //写入源文件长度
    to.write((const char *)count, sizeof(ULL) * 256); //写入频率表
    if (!to)
    {
        std::cerr << "[ERROR]: Failed to write to file `" << toName << "`." << std::endl;
        exit(6);
    }
    UB rem = 0;
    int remlen = 0;
    from.clear();
    from.seekg(from.beg);
    for (ULL j = 1; j <= flen; j++)
    {
        if (!from.get(c))
            break;
        code &cd = encoding[(UB)c];
        for (int i = 0; i < cd.len - 1; i++)
        {
            rem |= (cd.buf[i] >> remlen);
            if (!to.put(rem))
                goto end;
            rem = (cd.buf[i] << (8 - remlen));
        }
        rem |= (cd.buf[cd.len - 1] >> remlen);
        remlen += (8 - cd.bleft);
        if (remlen >= 8)
        {
            remlen -= 8;
            if (!to.put(rem))
                goto end;
            rem = (cd.buf[cd.len - 1] << (8 - cd.bleft - remlen));
        }
        if (!(j % 0x100000))
            std::cerr << "\rZipping: " << (j / 0x100000) << '/' << (flen / 0x100000) << "MB (" << (j * 100 / flen) << "%)";
    }
    std::cerr << "\rZipping: " << (flen / 0x100000) << '/' << (flen / 0x100000) << "MB (" << 100 << "%)" << std::endl;
    if (remlen)
        to.put(rem);
    end:
    if (!from)
    {
        std::cerr << "[ERROR]: Failed to read from file `" << fromName << "`." << std::endl;
        exit(3);
    }
    if (!to)
    {
        std::cerr << "[ERROR]: Failed to write to file `" << toName << "`." << std::endl;
        exit(6);
    }
    std::cerr << "Zipping finished. (" << flen << "B -> " << to.tellp() << "B)" << std::endl;

    //释放内存并关闭文件
    for (int i = 0; i < 256; i++)
        delete[] encoding[i].buf;
    from.close();
    to.close();
}

void unzip() //解压
{
    char buf[FILE_HEAD_LEN];
    ULL flen;

    //打开压缩文件
    std::ifstream from(fromName, std::ios::binary);
    if (!from)
    {
        std::cerr << "[ERROR]: Cannot open source file `" << fromName << "`." << std::endl;
        exit(2);
    }

    //检查文件头
    if (!from.read(buf, FILE_HEAD_LEN))
    {
        std::cerr << "[ERROR]: Failed to read from file `" << fromName << "`." << std::endl;
        exit(3);
    }
    if (strncmp(buf, FILE_HEAD, FILE_HEAD_LEN))
    {
        std::cerr << "[ERROR]: Incorrect formation of source file `" << fromName << "`." << std::endl;
        exit(4);
    }

    //读取原始文件长度
    from.read((char *)&flen, sizeof(ULL));
    if (!from)
    {
        std::cerr << "[ERROR]: Failed to read from file `" << fromName << "`." << std::endl;
        exit(3);
    }
    if (!flen) //原始文件为空
    {
        std::ofstream to(toName, std::ios::binary); //打开输出文件
        if (!to)
        {
            std::cerr << "[ERROR]: Cannot open output file `" << toName << "`." << std::endl;
            exit(5);
        }
        //关闭文件
        from.close();
        to.close();
        return;
    }

    //读取频率表
    from.read((char *)count, sizeof(ULL) * 256);
    if (!from)
    {
        std::cerr << "[ERROR]: Failed to read from file `" << fromName << "`." << std::endl;
        exit(3);
    }

    //根据频率表构建哈夫曼树
    tree ht = buildTree();

    //打开输出文件
    std::ofstream to(toName, std::ios::binary);
    if (!to)
    {
        std::cerr << "[ERROR]: Cannot open output file `" << toName << "`." << std::endl;
        exit(5);
    }

    //翻译压缩文件
    int p = ht.root;
    ULL j = 1;
    while (j <= flen)
    {
        if (!from.get(buf[0]))
        {
            if (from.eof())
            {
                std::cerr << "[ERROR]: Incorrect formation of source file `" << fromName << "`." << std::endl;
                exit(4);
            }
            std::cerr << "[ERROR]: Failed to read from file `" << fromName << "`." << std::endl;
            exit(3);
        }
        for (int i = 7; i >= 0; i--)
        {
            p = ((buf[0] >> i) & 1) ? nodes[p].right : nodes[p].left;
            if (!nodes[p].left)
            {
                if (!to.put(nodes[p].key))
                {
                    std::cerr << "[ERROR]: Failed to write to file `" << toName << "`." << std::endl;
                    exit(6);
                }
                if (++j > flen)
                    break;
                p = ht.root;
            }
        }
        if (!(j % 0x100000))
            std::cerr << "\rUnzipping: " << (j / 0x100000) << '/' << (flen / 0x100000) << "MB (" << (j * 100 / flen) << "%)";
    }
    std::cerr << "\rUzipping: " << (flen / 0x100000) << '/' << (flen / 0x100000) << "MB (" << 100 << "%)" << std::endl;

    //关闭文件
    from.close();
    to.close();
}

const char *USAGE = "usage: hzip <file name>";
int main(int argc, char **argv)
{
    if (argc != 2)
    {
        std::cerr << USAGE << std::endl;
        exit(1);
    }
    fromName = argv[1];
    if (fromName.length() > SUFFIX_LEN && fromName.substr(fromName.length() - SUFFIX_LEN, SUFFIX_LEN) == SUFFIX)
    {
        toName = fromName.substr(0, fromName.length() - SUFFIX_LEN);
        std::cerr << "Unzipping to file `" << toName << "`...." << std::endl;
        unzip();
        std::cerr << "Unzipped to file `" << toName << "`." << std::endl;
    }
    else
    {
        toName = fromName + SUFFIX;
        std::cerr << "Zipping to file `" << toName << "`...." << std::endl;
        zip();
        std::cerr << "Zipped to file `" << toName << "`." << std::endl;
    }
    return 0;
}
